import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Random;
import java.util.Vector;
import org.htmlparser.util.ParserException;

public class ScriptsDB   {

	public enum AppMode {LEARN, DETECT_FP, DETECT_FN}
	public static String newline = System.getProperty("line.separator");

	private HashMap<Integer,ScriptNode> appScripts;
	Vector<Integer> attackScriptsList =null; //this list holds all alternate hash values of attack scripts
	HashMap<Integer,String> learnedURLs = null; //this list contains all URL's learned during the learning phase
	private int globalCounter = 0;
	private BufferedWriter out;
	private Log detectorLog = null;
	private DetectorResults detectorResults = null;

	/**
	 * main constructor
	 * @param mode
	 * @throws DetectorException
	 * @throws IOException
	 */
	public ScriptsDB(String mode) throws DetectorException, IOException {
		//initialize lists
		attackScriptsList = new Vector<Integer>();
		attackScriptsList.clear();
		this.appScripts = new HashMap<Integer,ScriptNode>();
		appScripts.clear();
		this.learnedURLs = new HashMap<Integer,String>();
		learnedURLs.clear();
		this.detectorResults = new DetectorResults();

		if (mode.equals(AppMode.LEARN.toString()))
			out = new BufferedWriter(new FileWriter("LrnDB.txt"));
		else
			out = new BufferedWriter(new FileWriter("DetDB.txt"));
		this.detectorLog = Log.getInstance("detectorLog.txt");
	}

	public BufferedWriter learn(ArrayList<ExtendedLink> links, double learnPart) 
	throws IOException, DetectorException, ParserException{

		String hostPath = null;
		detectorLog.WriteToLog("Starting Learning Phase. The detector will learn " + 
				(Math.ceil((double)learnPart*links.size())) + "(" + learnPart*100 + "%) of total " + 
				links.size() +" urls",true);

		int learnCounter = 0;
		ExtendedLink exLink = null;
		String currentUrl = null;
		Vector<ScriptNode> pageScripts = null;


		//choose random url from list
		for (Iterator<ExtendedLink> iter = new RandomIterator<ExtendedLink>(links,new Random()); iter.hasNext();){
			learnCounter+=1;
			exLink = iter.next();
			currentUrl = exLink.getLink().toURL();
			if (currentUrl.indexOf('?') > -1)
				hostPath = currentUrl.substring(0, currentUrl.indexOf('?'));
			else
				hostPath = currentUrl;
			//add current URL hash to URL list
			learnedURLs.put(hostPath.hashCode(),currentUrl);

			//logging
			System.out.println("learning: " + currentUrl);
			detectorLog.WriteToLog("Checking for new scripts in " + currentUrl,false);
			out.write(newline + "url: " + currentUrl +newline + "--------------------------------------------" + newline + newline);

			ScriptExtractor2 scrext = null;
			try{
				scrext = new ScriptExtractor2(currentUrl);
			}catch (ParserException e){
				throw (new DetectorException(e.getMessage()));
			}
			
			//extract page script nodes
			pageScripts = scrext.getScriptNodes();
			
			//check if these scripts already exist in DB
			int[] newScriptsAdded = CheckScripts(pageScripts,AppMode.LEARN);
			if (newScriptsAdded[1]==0){
				System.out.println("No new scripts were found");
				detectorLog.WriteToLog("No new scripts were found",false);
			}
			else{
				System.out.println(newScriptsAdded[1] + " new scripts were found and added to scripts DB");
				detectorLog.WriteToLog(newScriptsAdded[1] + " new scripts were found and added to scripts DB",false);
			}
			if (learnCounter>=Math.ceil((double)learnPart*links.size()))
				break;
		}
		System.out.println("Learn phase finished. ");
		detectorLog.WriteToLog("Learn phase finished. " + newline,true);
		detectorLog.WriteToLog("The DB holds " + appScripts.size() + " scripts",false);
		return out;
	}

	public void setLearnedURLs(HashMap<Integer, String> learnedURLs) {
		this.learnedURLs = learnedURLs;
	}

	public HashMap<Integer, String> getLearnedURLs() {
		return learnedURLs;
	}

	public DetectorResults detect(ArrayList<ScriptNode> scriptsArray, ArrayList<ExtendedLink> links, double p, String mode) 
	throws IOException, ParserException, DetectorException{

		int detectCounter=0,alternateHashCounter=0,alternateHashAttack=0,effectiveUrlTested=0;
		int responsecode=0,overlap=0,normalHashCounter=0,normalHashAttack=0;
		int newURLAttackCounter=0;
		String hostPath=null;
		AppMode detectMode = AppMode.DETECT_FP;
		boolean urlIsEffectiveAttack=false;

		if (mode.equals("fn"))
			detectMode = AppMode.DETECT_FN;
		detectorLog.WriteToLog("Starting Detection  Phase. The detector will detect " + 
				(Math.ceil((double)p*links.size())) + "(" + p*100 + "%) of total " + 
				links.size() +" urls",true);
		Vector<ScriptNode> pageScripts = null;
		appScripts.clear();
		//convert scriptsArray to hashmap representation so we can detect script existence in O(1)
		for (ScriptNode s : scriptsArray){
			appScripts.put(s.getHash(), s);
		}
		//randomly traverse given links
		for (Iterator<ExtendedLink> iter = new RandomIterator<ExtendedLink>(links,new Random()); iter.hasNext();){
			urlIsEffectiveAttack=false;
			//increase number of examined URL's 
			detectCounter+=1;
			effectiveUrlTested+=1;
			//check if we reached max number of URL's 
			if (detectCounter>Math.ceil((double)p*links.size()))
				break;
			ExtendedLink exLink = iter.next();
			String currentUrl = exLink.getLink().toURL();
			currentUrl = currentUrl.replaceAll("\\s", "%20");
			if (currentUrl.indexOf('?') > -1)
				hostPath = currentUrl.substring(0, currentUrl.indexOf('?'));
			else
				hostPath = currentUrl;
			//if this URL has been learned, increase overlap factor
			if (learnedURLs.containsKey(hostPath.hashCode())) 
				overlap+=1; 
			System.out.println(newline + "examining: " + currentUrl);
			detectorLog.WriteToLog(detectCounter + ") Detecting for new scripts in " + currentUrl,false);
			out.write(newline + "url: " + currentUrl +newline + "--------------------------------------------" + newline + newline);
			//check if this URL is valid (i.e. , contains valid characters)
			//check if this is an SSL connection, as this might pose a problem
			if (currentUrl.startsWith("https")){
				System.out.println("This is a secured connection (HTTPS). Trying to reach the same URL over HTTP");
				currentUrl = currentUrl.replaceFirst("https", "http");
			}
			HttpURLConnection con = ((HttpURLConnection)(new URL(currentUrl)).openConnection());
			responsecode = con.getResponseCode();
			if (responsecode<200 || responsecode>=300){
				System.out.println("Server responded with " + responsecode + " " + con.getResponseMessage() + 
				". Skipping this URL");
				detectorLog.WriteToLog("Server responded with " + responsecode + " " + con.getResponseMessage() + 
						". Skipping this URL", false);
				effectiveUrlTested-=1;
				continue;
			}
			//extract page script nodes
			//pageScripts = new ScriptExtractor(currentUrl).getScriptNodes();
			pageScripts = new ScriptExtractor2(currentUrl).getScriptNodes();
			//check if these are new scripts or not
			int []newScriptsAdded = CheckScripts(pageScripts,detectMode);
			//first check the attacks that had been detecetd if not using canoninized form.
			if (newScriptsAdded[0]!=0){
				System.out.println("Detector found " + newScriptsAdded[0] + " scripts with different script code, " +
				"but their canoninized version was the same");
				//increment number of URL's that might cause an alarm if we didn't use alternate Hash
				normalHashCounter+=1;
				normalHashAttack+=newScriptsAdded[0];
				//decrement number of effective attacks
			}
			//now check real attacks
			if (newScriptsAdded[1]==0){
				System.out.println("No FP/FN scripts were found in this URL");
				detectorLog.WriteToLog("No FP/FN scripts were found",false);
			}
			else{//found new script nodes
				//check if this URL has been learned
				if (!learnedURLs.containsKey(hostPath.hashCode())){
					newURLAttackCounter+=1;
					System.out.println("Found " + newScriptsAdded[1] +" new scripts, but this is URL hasn't been learned " +
					" --> attack doesn't count");
					detectorLog.WriteToLog("Found " + newScriptsAdded[1] +" new scripts, but this is URL hasn't been learned " +
							" --> attack doesn't count",false);
				}
				else{//this URL has been learned
					detectorLog.WriteToLog(newScriptsAdded[1] + " new scripts were found  - ATTACK",false);
					alternateHashCounter+=1;
					alternateHashAttack+=newScriptsAdded[1];
					urlIsEffectiveAttack=true;
				}
			}
			if (urlIsEffectiveAttack==false)
				effectiveUrlTested-=1;
		}
		out.close();
		detectorResults.setEffectiveUrlTested(effectiveUrlTested);
		detectorResults.setNumNormalizedAttackScripts(alternateHashAttack);
		detectorResults.setNumNormalizedUrlWithAttackScripts(alternateHashCounter);
		detectorResults.setNumRegularUrlWithAttackScripts(normalHashCounter);
		detectorResults.setNumUnNormalizedAttackScripts(normalHashAttack);
		detectorResults.setOverlappingFactor(overlap);
		detectorResults.setNumUnLearnedUrlsWithAttack(newURLAttackCounter);

		return detectorResults;
	}
/**
 * 
 * @param pageScripts
 * @param mode
 * @return
 * @throws IOException
 * @throws DetectorException
 */
	public int[] CheckScripts(Vector<ScriptNode> pageScripts, AppMode mode) throws IOException, DetectorException{

		ScriptNode currentNode;
		int currentAlternativeHashVal;
		boolean match = false, xssAttackIndicator=false;
		
		// localCounter[0] - number of scripts with different script code but similar canoninized version.
		// localCounter[1] - number of scripts with no match at all.
		int[] localCounters = {0,0}; 

		if (pageScripts!=null){
			//need to go over all scripts and check for each script if it is in internal scripts DB
			for (Iterator<ScriptNode> iter = pageScripts.iterator();iter.hasNext();){ 

				match=false;
				xssAttackIndicator=false;

				currentNode = (ScriptNode)iter.next(); // current script
				
				//check if this script is an injected attack
				if ((currentNode.getOrigCode().toLowerCase().indexOf("xss")>-1) || 
						(currentNode.getOrigCode().toLowerCase().indexOf("fromcharcode")>-1))
					xssAttackIndicator = true;
				
				//check if learnDB already contains script
				
				//first check native hash and then check alternative hash
				if (!appScripts.containsKey(currentNode.getHash())){ 
					// the script is not in DB according to regular hash
					currentAlternativeHashVal = currentNode.getAlternateHash();
					// go over all scripts in DB according to alternative hash value
					Collection<ScriptNode> col = appScripts.values(); // get all scripts from DB
					Iterator<ScriptNode> i = col.iterator(); 
					if  ( (currentNode.getAttr() != null) &&
							(!currentNode.getOrigCode().endsWith(".js")) &&
							(!(currentNode.getScriptType()==ScriptNode.ScriptType.URI_SCHEME))
					){ // this node's source is not JavaScript. ignore it.
						continue;
					}
					while (i.hasNext()){
						if (i.next().getAlternateHash()==currentAlternativeHashVal){ 
							match = true;
							/* found a match - this means that although this JS doesn't have
							 * the same code as other script in script DB, it still has
							 * the same canoninized code.
							 * Still, we want to log this event 
							 */
							//check if this script has been seen before
							if (!attackScriptsList.contains(currentAlternativeHashVal)){
								localCounters[0]+=1;
								globalCounter+=1;
								attackScriptsList.add(currentAlternativeHashVal);
								out.write("-----------------------------------------------------------------" + newline +
										"Normal Hash Script # " + globalCounter + newline +
										"-----------------------------------------------------------------" + newline + 
										currentNode.getOrigCode() + newline);
							}
							break;
						}
					}
				}
				else{
					match=true;
				}
				if (match==false){
					// no match (regular/alternate) was found in scriptDB
					if (mode==AppMode.LEARN){
						//learn new script
						globalCounter+=1;
						localCounters[1]+=1;
						//add new script to scriptDB
						appScripts.put(currentNode.getHash(), currentNode);
						//System.out.println("New script was found and added" + newline);
						out.write("-----------------------------------------------------------------" + newline +
								"Script # " + globalCounter + newline +
								"-----------------------------------------------------------------" + newline + 
								newline + currentNode.getOrigCode() + newline);
					}
					else if (mode==AppMode.DETECT_FP){//DETECT mode
						//check if new script hasn't been detected before
						if (!attackScriptsList.contains(currentNode.getAlternateHash())){
							globalCounter+=1;
							localCounters[1]+=1;
							out.write("-----------------------------------------------------------------" + newline +
									"False Positive Script #"  + globalCounter + newline +
									"-----------------------------------------------------------------" + newline + 
									currentNode.getScriptType() + newline + 
									currentNode.getOrigCode() + newline);
							attackScriptsList.add(currentNode.getAlternateHash());

						}
						else {
							detectorLog.WriteToLog("Found a false positive script which is in the AttackDB.",false);
							detectorLog.WriteToLog("Script details:",false);
							detectorLog.WriteToLog(currentNode.getOrigCode(),false);
						}
					}
					else if (mode==AppMode.DETECT_FN){
						//check if this is an injected attack script
						if (xssAttackIndicator==true)
							System.out.println("Detected an injected attack");
						else
							System.out.println("Detected an attack script that is not an injected attack");

						//check if this new script hasn't been detected before
						//if (!attackScriptsList.contains(currentNode.getAlternateHash())){
						globalCounter+=1;
						localCounters[1]+=1;
						out.write("-----------------------------------------------------------------" + newline +
								"Attack Script #"  + globalCounter + newline +
								"-----------------------------------------------------------------" + newline + 
								currentNode.getScriptType() + newline + 
								currentNode.getOrigCode() + newline);
						attackScriptsList.add(currentNode.getAlternateHash());

						//}
						//else {
						//detectorLog.WriteToLog("Found an attack script which is already added to AttackDB.",false);
						//detectorLog.WriteToLog("Script details:",false);
						//detectorLog.WriteToLog(currentNode.getScriptCode(),false);
						//}
					}
				}
			}
		}
		return (localCounters);
	}


	public ArrayList<ScriptNode> getAppScripts() {
		ArrayList<ScriptNode> arr = new ArrayList<ScriptNode>();
		Collection<ScriptNode> col = this.appScripts.values();
		for (ScriptNode s : col){
			arr.add(s);
		}
		return arr;
	}


}
